#!/usr/bin/env python

import time
import sys
import os
import subprocess
import logging
import paramiko, socket
import shlex
import argparse
import ConfigParser
import re
from mailer import Mailer
from mailer import Message


class CheckUsbError(Exception):
    pass

class SshExitNonZero(Exception):
    pass


def sync_buffers():
    """Force changed blocks to disk, log and pass exceptions."""
    try:
        subprocess.check_call(sync)
    except (OSError, subprocess.CalledProcessError) as e:
        logging.error(e)
        pass


def mount_usb(backup_disk):
    """Mount backup_disk and raise exception on failure."""
    try:
        subprocess.check_call([mount, backup_disk])
        logging.debug('Mounted %s successfully.', backup_disk)
    except (OSError, subprocess.CalledProcessError) as e:
        try:
            raise
        finally:
            pass


def umount_usb(backup_disk):
    """Unmount backup backup_disk and raise exception on failure."""
    try:
        sync_buffers()
        subprocess.check_call([umount, backup_disk])
        logging.debug('Unmounted %s successfully.', backup_disk)
    except (OSError, subprocess.CalledProcessError) as e:
        try:
            raise
        finally:
            pass


def check_usb(backup_disk, basedir):
    """Test basedir exists and is rw on backup_disk."""
    try:
        if not os.path.isdir(os.path.join(backup_disk, basedir)):
            raise CheckUsbError('%s not found in %s.' %(basedir, backup_disk))
        if not os.access(os.path.join(backup_disk, basedir), os.W_OK|os.R_OK):
            raise CheckUsbError('%s is not writeable.' %(
                            os.path.join(backup_disk, basedir)))
        logging.debug('Backup disk validation passed.')
    except (OSError, CheckUsbError) as e:
        try:
            raise
        finally:
            pass


def ssh_cmd(host, user, cmd):
    """
    SSH to vba host, return exit status, stdout and stderr of cmd.
    Depends on ssh shared keys. Has some leniency for socket timeouts.
    """
    try:
        client = paramiko.SSHClient()
        client.load_system_host_keys()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        socket_timeout_retries = range(3)
        first_error = None
        for attempt in socket_timeout_retries:
            try:
                client.connect(host, username=user, timeout=30)
                first_error = None
                break
            except (socket.timeout) as st:
                logging.warning(st)
                attempt + 1
                if not first_error:
                    first_error = st
                time.sleep(30)
        if first_error:
            try:
                logging.error('Failed to connect to \"%s\"', host)
                raise first_error
            finally:
                pass

        channel = client.get_transport().open_session()
        channel.exec_command(cmd)
        exitstatus = channel.recv_exit_status()
        stdout = channel.recv('1024')
        stderr = channel.recv_stderr('1024')
        channel.close()
        return exitstatus, stdout, stderr
    except (paramiko.SSHException, socket.error) as e:
        try:
            raise
        finally:
            pass


def vba_jobs_running(host, user):
    """
    Connect to vba host and check for running backups. If backups or
    Post-Process queue jobs are running returns True else returns False.
    If exception is raised job stops processing.
    """
    try:
        jobs_running = False
        cmd = '/usr/bin/screendump'
        exitstatus, stdout, stderr = ssh_cmd(host, user, cmd)
        if exitstatus != 0:
            raise SshExitNonZero(stderr.decode('utf-8'))
        logging.debug('Obtain vba job status succeeded')
        w = '.Worker-\d: '
        p = '.*Post-Process Queue: [0-9]'
        for line in stdout.splitlines():
            if re.match(w, line):
                if line.split()[1] != '(idle)':
                    logging.debug('Workers not idle')
                    jobs_running = True
                    break
                else:
                    logging.debug('Workers all idle')

            if re.match(p, line):
                n = int(line.split()[line.split().index('Post-Process')+2])
                if n != 0:
                    logging.debug('Post-Process Queue still running')
                    jobs_running = True
                    break
                else:
                    logging.debug('Post-Process Queue idle')

        return jobs_running
    except (SshExitNonZero, paramiko.SSHException, socket.error) as e:
        try:
            raise
        finally:
            pass


def sync_vba(options, exclude, src, dest):
    output = None
    try:
        o = shlex.split(options)
        if exclude:
            e = ['--exclude=%s' % x.strip() for x in exclude.split(',')]
            r = ([rsync] + o + e + [src, dest])
        else:
            r = ([rsync] + o + [src, dest])
        logging.debug('Rsync command: %s', r)
        output = subprocess.check_output(r, stderr=subprocess.STDOUT)
        logging.info('Rsync of %s completed successfully', src)
        logging.info(output.decode('utf-8'))
        return output.decode('utf-8')
    except Exception as e:
        try:
            raise
        finally:
            try:
                logging.error('Rsync of %s encountered and error', src)
                logging.error(output.decode('utf-8'))
                # i think this is wrong and needs testing.
                return output.decode('utf-8')
            except:
                pass


def process_jobs(vba_jobs_iteritems):
    """Loop through jobs."""
    all_jobs_success = True
    all_jobs_start = time.time()
    body = []
    # alter to some blah de blah message intro
    body.append("Began running at: %s\n\n" % time.strftime('%c %Z', time.localtime(all_jobs_start)))
    for key, value in vba_jobs_iteritems:
        try:
            job_success = False
            mounted_usb_ok = False
            rsyncout = 'No Output logged\n'
            job_start = time.time()
            rsync_start = job_start
            timewindow = []

            jobdir = key
            host = value.get('vbahost')
            user = value.setdefault('user', 'root')
            # need to correctly set srcdir to ensure trailing "/"
            srcdir = value.setdefault('srcdir', '/dd/')
            exclude = value.setdefault('exclude', None)
            localpath = value.setdefault('localpath', None)
            basedir = value.setdefault('basedir', 'VBABACKUPS')
            dontstartbefore = value.setdefault('dontstartbefore', 18)
            dontstartafter = value.setdefault('dontstartafter', 6)
            backup_disk = value.setdefault('usbmount', '/mnt/backup')
            backup_disk_mounted = os.path.ismount(backup_disk)

            if force:
                timewindow = range(0,24)
            else:
                st = int(dontstartbefore)
                et = int(dontstartafter)
                if st <= et:
                    timewindow = range(st, et)
                else:
                    timewindow = range(st,24) + range(0,et) 

            # Look at adding StrictHostKeyChecking=no as per:
            # http://askubuntu.com/questions/123072/ssh-automatically-accept-keys
            # this needs to go into a function...  and use a var for rsync options.
            # also option to use vanilla rsync if srcdir startswith ':'
            if dry_run:
                options = value.setdefault('rsync_options', '-ahHAX -e "ssh -c arcfour" --stats --delete --whole-file') + ' --dry-run'
            else:
                options = value.setdefault('rsync_options', '-ahHAX -e "ssh -c arcfour" --stats --delete --whole-file')

            vbachk_cmd = 'ls ' + srcdir + 'phdvbstore.ini > /dev/null 2>&1'
            if ssh_cmd(host, user, vbachk_cmd)[0] != 0:
                raise Exception('%s host connect and validate failed for %s' %(host, vbachk_cmd))

            dest = os.path.join(backup_disk, basedir, jobdir)
            if localpath:
                d = localpath + '/'
                if not os.path.isfile(d + 'phdvbstore.ini'):
                    raise Exception('localpath %s is invalid.' %(d))
                else:
                    src = d
            else:
                src = user + '@' + host + ':' + srcdir

            """Mount backup disk if not already mounted."""
            # should probably look at putting this after timewindow checks
            if not backup_disk_mounted:
                logging.debug('Attempting to mount %s', backup_disk)
                mount_usb(backup_disk)
                mounted_usb_ok = True
            else:
                logging.debug('%s already mounted', backup_disk)

            """Check the backup disk"""
            check_usb(backup_disk, basedir)

            # need to reimplement scheduling
            while True:
                if time.localtime()[3] in timewindow:
                    if vba_jobs_running(host, user):
                        zz = 300
                        logging.debug('%s still running, sleeping %s seconds', jobdir, zz)
                        time.sleep(zz)
                    else:
                        rsync_start = time.time()
                        rsyncout = sync_vba(options, exclude, src, dest)
                        job_success = True
                        break
                else:
                    # shouldn't i just raise and exception here?
                    # it would work the same as break
                    job_success = False
                    all_jobs_success = False
                    logging.info('Job outside of timewindow')
                    body.append('!!! JOB OUTSIDE OF ALLOCATED TIME WINDOW !!! : %s\n\n' % timewindow)
                    break

        except Exception as e:
            logging.exception(e)
            job_success = False
            all_jobs_success = False

        finally:
            job_end = time.time()
            rsync_end = job_end
            if mounted_usb_ok:
                if not backup_disk_mounted:
                    logging.debug('Attempting to unmount %s', backup_disk)
                    umount_usb(backup_disk)
                else:
                    logging.debug('%s left mounted', backup_disk)

            body.append(report_generator(jobdir, job_success, job_start, job_end,
                                         rsync_start, rsync_end, rsyncout))

    all_jobs_end = time.time()
    body.append("Processing finshed at: %s\n" % time.strftime('%c %Z', time.localtime(all_jobs_end)))
    return all_jobs_success, body 
             

def relay_email(relayhost, relayauth, subject, body,
                recipient, sender_email, log_file):
    b = ''.join(body)
    try:
        message = Message(From=sender_email,
                          To=recipient,
                          Subject=subject)
        message.Body = b
        message.attach(log_file)
        sender = Mailer(relayhost)
        sender.send(message)
    except Exception as e:
        try:
            logging.error(e)
        finally:
            pass


def report_generator(jobdir, job_success, job_start, job_end, rsync_start, rsync_end, rsyncout):
    s = []
    s.append('------ %s ------\n' % jobdir)
    s.append('Job success: %s\n\n' % job_success)
    s.append('------ TIME ------\n')
    s.append("Started at: %s\n" % time.strftime('%c %Z', time.localtime(job_start)))
    s.append("Ended at: %s\n" % time.strftime('%c %Z', time.localtime(job_end)))
    s.append("Duration: %s minutes\n\n" % int((job_end - job_start) / 60))
    s.append("------ JOB STATISTICS ------\n")
    s.append("Rsync Started at: %s\n" % time.strftime('%c %Z', time.localtime(rsync_start)))
    s.append("Rsync Ended at: %s\n" % time.strftime('%c %Z', time.localtime(rsync_end)))
    s.append("Duration: %s minutes\n" % int((rsync_end - rsync_start) / 60))
    s.append(rsyncout)
    summary = ''.join(s)
    return summary


def main():

    log_file = '/var/log/syncvba.log'
    exit_status = 1 
    body = None
    all_jobs_result = None

    """Setup argparse."""
    parser = argparse.ArgumentParser()

    parser.add_argument('-c', '--config-file',
            help='path to config file, default = /etc/syncvba/syncvba.ini',
            default='/etc/syncvba/syncvba.ini')

    parser.add_argument('-l', '--log-level',
            help='set logging level',
            default='INFO',
            choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'])

    parser.add_argument('-d', '--dry-run',
            help='Perform dry run',
            action='store_true')

    parser.add_argument('-x', '--cleanup',
            help='Clean up usb mounts.',
            action='store_true')

    parser.add_argument('-f', '--force',
            help='Run backup outside of time window',
            action='store_true')

    parser.add_argument('-n', '--no-email',
            help='Disable email reports',
            action='store_true')

    """Create variables from argparse."""
    args = parser.parse_args()
    log_level = args.log_level
    config_file = args.config_file
    global dry_run
    dry_run = args.dry_run
    cleanup = args.cleanup
    global force
    force = args.force
    no_email = args.no_email

    # Implement log rotation from section "15.6.1.1. Simple examples"
    # http://docs.activestate.com/activepython/3.1/python/library/logging.html
    logging.basicConfig(
            filename=log_file, format='%(asctime)s %(levelname)s: %(message)s',
            datefmt='%F %H:%M:%S', filemode='w',
            level=log_level)

    config = ConfigParser.ConfigParser()
    try:
        if not os.path.isfile(config_file):
            raise Exception('Cannot find config file: %s' %(config_file))
        config.read(config_file)
    except Exception as e:
        logging.error(e)
        return 1

    """Obtain general, default and job variables."""
    try:
        general = dict(config.items('GENERAL'))
        if no_email:
            relayhost = False
        else:
            relayhost = general.setdefault('smtp_relay', False)
        sender_email = general.setdefault('smtp_sender', None)
        recipient = general.setdefault('smtp_recipients', None).split()
        if dry_run:
            subject = general.setdefault('smtp_subject', 'VBA Sync') + ' !!DRY RUN!!'
        else:
            subject = general.setdefault('smtp_subject', 'VBA Sync')
        relayauth = general.setdefault('smtp_auth', None)

        global rsync, sync, mount, umount
        rsync = general.setdefault('rsync_path', '/usr/bin/rsync')
        sync = general.setdefault('sync_path', '/bin/sync')
        mount = general.setdefault('mount_path', '/bin/mount')
        umount = general.setdefault('umount_path', '/bin/umount')

        vba_sections = [x for x in config.sections() if x not in ['GENERAL']]
        vba_jobs = {}
        for section in vba_sections:
            vba_jobs[section] = dict(config.items(section))
        xvba_jobs = len(vba_jobs)
        logging.debug('Found %s backup jobs: %s', xvba_jobs, vba_sections)
    except Exception as e:
        logging.error(e)
        return 1

    all_jobs_result, body = process_jobs(vba_jobs.iteritems())
    if all_jobs_result:
        exit_status = 0
        subject = subject + ' SUCCESS'
    else:
        subject = subject + ' FAILURE'

    if relayhost:
        try:
            relay_email(relayhost, relayauth, subject, body,
                        recipient, sender_email, log_file)
            logging.debug('Email sent successfully')
        except Exception as e:
            logging.error(e)
            exit_status = 1

    return exit_status


if __name__ == '__main__':
    sys.exit(main())

